//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package impl

import (
	"context"
	"sync"

	"atomgit.com/openlinksaas/ssh-proxy/gen_code"
	"github.com/google/uuid"
	"github.com/pkg/sftp"
)

type ConnectionApiImpl struct{}

func (s *ConnectionApiImpl) ConnectionClosePost(ctx context.Context, req gen_code.ConnectionClosePostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	if client.sftpClient != nil {
		client.sftpClient.Close()
	}

	client.forwardMap.Range(func(_, value any) bool {
		value.(*ForwardClient).Close()
		return true
	})
	client.shellMap.Range(func(_, value any) bool {
		info := value.(*ShellInfo)
		if info.stdin != nil {
			info.stdin.Close()
		}
		if info.session != nil {
			info.session.Close()
		}
		close(info.recvChan)
		return true
	})
	client.client.Close()
	clientMap.Delete(req.Id)
	return gen_code.ImplResponse{
		Code: 200,
		Body: req.Id,
	}, nil
}

func (s *ConnectionApiImpl) ConnectionListPost(ctx context.Context, req gen_code.ConnectionListPostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	clientList := []gen_code.ConnectionListPost200ResponseInner{}

	clientMap.Range(func(k, v any) bool {
		info := v.(*ClientInfo)
		authList := []gen_code.ConnectionListPost200ResponseInnerAuthListInner{}
		for _, auth := range info.AuthInfoList {
			authList = append(authList, gen_code.ConnectionListPost200ResponseInnerAuthListInner{
				Username: auth.Username,
				Addr:     auth.Addr,
			})
		}
		clientList = append(clientList, gen_code.ConnectionListPost200ResponseInner{
			Id:       k.(string),
			Name:     info.Name,
			AuthList: authList,
		})
		return true
	})
	return gen_code.ImplResponse{
		Code: 200,
		Body: clientList,
	}, nil
}

func (s *ConnectionApiImpl) ConnectionOpenPost(ctx context.Context, req gen_code.ConnectionOpenPostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	var parent *ChainClient
	for _, authInfo := range req.AuthInfoList {
		client, err := NewChainClient(parent, &authInfo)
		if err != nil {
			if parent != nil {
				parent.Close()
			}
			return gen_code.ImplResponse{
				Code: 500,
				Body: &gen_code.ErrInfo{
					ErrMsg: err.Error(),
				},
			}, nil
		}
		parent = client
	}
	sftpClient, _ := sftp.NewClient(parent.realClient) //skip error check
	id := uuid.NewString()
	authList := []UsernameAndAddr{}
	for _, auth := range req.AuthInfoList {
		authList = append(authList, UsernameAndAddr{
			Username: auth.Username,
			Addr:     auth.Addr,
		})
	}
	clientMap.Store(id, &ClientInfo{
		Name:         req.Name,
		AuthInfoList: authList,
		client:       parent,
		sftpClient:   sftpClient,
		forwardMap:   sync.Map{},
		shellMap:     sync.Map{},
	})
	return gen_code.ImplResponse{
		Code: 200,
		Body: id,
	}, nil
}
